# 数据库模型和操作
import sqlite3
import json
from datetime import datetime
from typing import List, Dict, Optional

class Database:
    def __init__(self, db_path: str = "toutiao_auto.db"):
        self.db_path = db_path
        self.init_database()
    
    def init_database(self):
        """初始化数据库表"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 账号表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS accounts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                nickname TEXT,
                avatar_url TEXT,
                cookies TEXT,
                is_active BOOLEAN DEFAULT 1,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 文章表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS articles (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                title TEXT NOT NULL,
                content TEXT NOT NULL,
                prompt TEXT,
                account_id INTEGER,
                status TEXT DEFAULT 'draft',
                article_url TEXT,
                views INTEGER DEFAULT 0,
                likes INTEGER DEFAULT 0,
                comments INTEGER DEFAULT 0,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                published_at TIMESTAMP,
                FOREIGN KEY (account_id) REFERENCES accounts (id)
            )
        ''')
        
        # 提示词模板表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS prompt_templates (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                category TEXT,
                prompt TEXT NOT NULL,
                description TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        conn.commit()
        conn.close()
    
    def save_account(self, username: str, nickname: str = None, 
                    avatar_url: str = None, cookies: str = None) -> int:
        """保存账号信息"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT OR REPLACE INTO accounts 
            (username, nickname, avatar_url, cookies, updated_at)
            VALUES (?, ?, ?, ?, ?)
        ''', (username, nickname, avatar_url, cookies, datetime.now()))
        
        account_id = cursor.lastrowid
        conn.commit()
        conn.close()
        return account_id
    
    def get_active_account(self) -> Optional[Dict]:
        """获取当前激活的账号"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT id, username, nickname, avatar_url, cookies
            FROM accounts 
            WHERE is_active = 1
            ORDER BY updated_at DESC
            LIMIT 1
        ''')
        
        result = cursor.fetchone()
        conn.close()
        
        if result:
            return {
                'id': result[0],
                'username': result[1],
                'nickname': result[2],
                'avatar_url': result[3],
                'cookies': result[4]
            }
        return None
    
    def save_article(self, title: str, content: str, prompt: str = None, 
                    account_id: int = None) -> int:
        """保存文章"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT INTO articles (title, content, prompt, account_id)
            VALUES (?, ?, ?, ?)
        ''', (title, content, prompt, account_id))
        
        article_id = cursor.lastrowid
        conn.commit()
        conn.close()
        return article_id
    
    def update_article_status(self, article_id: int, status: str, 
                            article_url: str = None):
        """更新文章状态"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        if status == 'published':
            cursor.execute('''
                UPDATE articles 
                SET status = ?, article_url = ?, published_at = ?
                WHERE id = ?
            ''', (status, article_url, datetime.now(), article_id))
        else:
            cursor.execute('''
                UPDATE articles SET status = ? WHERE id = ?
            ''', (status, article_id))
        
        conn.commit()
        conn.close()
    
    def get_articles(self, account_id: int = None, limit: int = 50) -> List[Dict]:
        """获取文章列表"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        if account_id:
            cursor.execute('''
                SELECT id, title, content, status, views, likes, comments, 
                       created_at, published_at
                FROM articles 
                WHERE account_id = ?
                ORDER BY created_at DESC
                LIMIT ?
            ''', (account_id, limit))
        else:
            cursor.execute('''
                SELECT id, title, content, status, views, likes, comments,
                       created_at, published_at
                FROM articles 
                ORDER BY created_at DESC
                LIMIT ?
            ''', (limit,))
        
        results = cursor.fetchall()
        conn.close()
        
        articles = []
        for row in results:
            articles.append({
                'id': row[0],
                'title': row[1],
                'content': row[2][:200] + '...' if len(row[2]) > 200 else row[2],
                'status': row[3],
                'views': row[4],
                'likes': row[5],
                'comments': row[6],
                'created_at': row[7],
                'published_at': row[8]
            })
        
        return articles
    
    def get_article(self, article_id: int) -> Optional[Dict]:
        """根据ID获取文章"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT id, title, content, status, article_url, views, likes, comments,
                   created_at, published_at
            FROM articles 
            WHERE id = ?
        ''', (article_id,))
        
        result = cursor.fetchone()
        conn.close()
        
        if result:
            return {
                'id': result[0],
                'title': result[1],
                'content': result[2],
                'status': result[3],
                'url': result[4],
                'views': result[5],
                'likes': result[6],
                'comments': result[7],
                'created_at': result[8],
                'published_at': result[9]
            }
        return None
    
    def update_article_stats(self, article_id: int, stats: Dict[str, int]):
        """更新文章统计数据"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            UPDATE articles 
            SET views = ?, likes = ?, comments = ?
            WHERE id = ?
        ''', (stats.get('views', 0), stats.get('likes', 0), 
              stats.get('comments', 0), article_id))
        
        conn.commit()
        conn.close()
    
    def save_prompt_template(self, name: str, category: str, 
                           prompt: str, description: str = None) -> int:
        """保存提示词模板"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT INTO prompt_templates (name, category, prompt, description)
            VALUES (?, ?, ?, ?)
        ''', (name, category, prompt, description))
        
        template_id = cursor.lastrowid
        conn.commit()
        conn.close()
        return template_id
    
    def get_prompt_templates(self, category: str = None) -> List[Dict]:
        """获取提示词模板"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        if category:
            cursor.execute('''
                SELECT id, name, category, prompt, description
                FROM prompt_templates 
                WHERE category = ?
                ORDER BY created_at DESC
            ''', (category,))
        else:
            cursor.execute('''
                SELECT id, name, category, prompt, description
                FROM prompt_templates 
                ORDER BY created_at DESC
            ''')
        
        results = cursor.fetchall()
        conn.close()
        
        templates = []
        for row in results:
            templates.append({
                'id': row[0],
                'name': row[1],
                'category': row[2],
                'prompt': row[3],
                'description': row[4]
            })
        
        return templates
